package com.gujiangbo.application.common.base;

import java.util.HashMap;
import java.util.Map;
import javax.annotation.PostConstruct;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.PropertySources;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Repository;

@Component
@PropertySources({
		@org.springframework.context.annotation.PropertySource({ "classpath:config/base/zookeeper.properties" }) })
@Repository("zkClientPool")
public class ZKClientPool {
	@Value("${zookeeper.url}")
	private String zookeeperUrl;
	@Value("${zookeeper.baseSleepTimeMs}")
	private int baseSleepTimeMs;
	@Value("${zookeeper.maxRetries}")
	private int maxRetries;
	@Value("${zookeeper.sessionTimeoutMs}")
	private int sessionTimeoutMs;
	private RetryPolicy retryPolicy;
	private static CuratorFrameworkFactory.Builder zkCliBuilder;
	private static Map<String, CuratorFramework> poolMap = new HashMap<String, CuratorFramework>();

	@PostConstruct
	private void init() {
		this.retryPolicy = new ExponentialBackoffRetry(this.baseSleepTimeMs, this.maxRetries);
		zkCliBuilder = CuratorFrameworkFactory.builder().connectString(this.zookeeperUrl)
				.sessionTimeoutMs(this.sessionTimeoutMs).retryPolicy(this.retryPolicy);
	}

	public CuratorFramework getClient(String key, String namespace) {
		CuratorFramework cli = (CuratorFramework) poolMap.get(key);
		if (cli == null) {
			cli = zkCliBuilder.namespace(namespace).build();
			poolMap.put(key, cli);
		}
		if (cli.getState().compareTo(CuratorFrameworkState.STARTED) != 0) {
			cli.start();
		}
		return cli;
	}

	public CuratorFramework getClient(String key) {
		return (CuratorFramework) poolMap.get(key);
	}

	public void close(String key) {
		CuratorFramework cli = (CuratorFramework) poolMap.get(key);
		if (cli == null) {
			return;
		}
		try {
			cli.close();
		} catch (Exception localException) {
		}
	}
}
